'''
Created on Jan 25, 2011

@author: dexen
'''

import os
import sys
import threading

import rpyc

from dexen.libs import ranking
from dexen.system.common import service, utils
from dexen.system.common.meta import ScoreMeta
from dexen.system.common.base import Action, Individual
from dexen.libs.ranking import ParetoRanking, RankingResult
from dexen.system.common.excepts import (JobNotExistsException,
                                         DexenInvalidJobDefError,
                                         DexenConnectionError,
                                         DexenInvalidJobNameError)    


#===============================================================================
# Helper Classes
#===============================================================================
class BaseProgress(object):
    def __init__(self, inds, scores, step):
        self.inds = inds # list of Individual
        self.scores = scores # list of ScoreMeta
        self.step = step # step size

    def _is_fully_evaluated(self, ind):
        for score in self.scores:
            assert isinstance(score, ScoreMeta)
            if not hasattr(ind, score.name):
                return False
            if getattr(ind, score.name) == None:
                return False
        return True
    
    def _get_live_inds(self, time):
        alive_inds = []
        for ind in self.inds:
            assert isinstance(ind, Individual)
            if ind._act_time > time: 
                continue
            if ind._inact_time != None and ind._inact_time <= time: 
                continue
            if self._is_fully_evaluated(ind):
                alive_inds.append(ind)
        return alive_inds
        

class ScoreProgress(BaseProgress):
    def __init__(self, inds, scores, step):
        BaseProgress.__init__(self, inds, scores, step)
        
    @classmethod
    def get_header(self, scores):
        header = "Time"
        for score in scores:
            assert isinstance(score, ScoreMeta)
            header += ",%s_best,%s_avg,%s_worst"%(score.name, score.name, 
                                                  score.name)
        return header    
                   
    def get_progress(self):
        result = []
        inds = self.inds
        inds.sort(key=lambda ind: ind._act_time)
        index = 1
        INF = sys.maxint
        while (index < len(inds)):
            ind = inds[index]
            assert isinstance(ind, Individual)
            time = ind._act_time
            #alive and fully evaluated inds
            alive_inds = self._get_live_inds(time) 

            score_info = {}
            for score in self.scores:
                assert isinstance(score, ScoreMeta)
                if score.is_maximize():
                    worst, best = INF, -INF
                else:
                    worst, best = -INF, INF
                total = 0.0
                for ind in alive_inds:
                    score_val = getattr(ind, score.name)
                    assert score_val != None
                    total += score_val 
                    if score.is_maximize():
                        if score_val > best:
                            best = score_val
                        if score_val < worst:
                            worst = score_val
                    else: 
                        if score_val < best:
                            best = score_val
                        if score_val > worst:
                            worst = score_val
                if len(alive_inds) == 0:
                    avg = 0.0
                else:
                    avg = total / len(alive_inds)
                score_info[score.name] = (best, avg, worst)
            result.append((time, score_info))
            index += self.step
        return result
    

class ParetoAdvancement(BaseProgress):
    def __init__(self, inds, scores, step):
        BaseProgress.__init__(self, inds, scores, step)
    
    @classmethod
    def get_header(cls, scores, adv_list):
        header = ""
        for time, _ in adv_list:
            header += "id@%d"%(time,)
            for score in scores:
                assert isinstance(score, ScoreMeta)
                header += ",%s@%d"%(score.name, time)
            header += ","
        return header[:-1]              
    
    def get_advancement(self):
        result = []
        inds = self.inds
        inds.sort(key=lambda ind: ind._act_time)
        index = 1
        while (index < len(inds)):
            ind = inds[index]
            assert isinstance(ind, Individual)
            time = ind._act_time
            #alive and fully evaluated inds
            alive_inds = self._get_live_inds(time) 
            p_ranking = ParetoRanking(alive_inds)
            p_ranking.set_scores(self.scores)
            ranked = p_ranking.rank(1) # returns list of RankingResult
            ranked_inds = [rank_result.ind for rank_result in ranked]
            ranked_inds.sort(key=lambda ind: ind.__id__)
            result.append((time, ranked_inds))
            index += self.step
        return result

            
#===============================================================================
# Server Proxy
#===============================================================================
class ServerProxy(object):
    def __init__(self):
        self.srv_conn = None
        self.job_proxies = {} # job_id : JobProxy
        self.lock = threading.RLock()
        self.svr_ip = None
        self.svr_start_time = None

    def connect(self, reg_host):
        self.srv_conn = service.get_svr_conn(reg_host)
        self.svr_ip = self.srv_conn.root.get_ip()
        self.svr_start_time = self.srv_conn.root.get_start_time()
        
        # Reconnect all the cached job_proxies
        for job_proxy in self.job_proxies.values():
            assert isinstance(job_proxy, JobProxy)
            job_proxy.reconnect(self.srv_conn)

    def disconnect(self):
        self.srv_conn.close()

    def is_connected(self):
        return self.srv_conn
    
    def check_connection(self):
        if self.srv_conn is None:
            err_msg = "Server is not reachable.\n"
            err_msg += "Please start a server first."
            raise DexenConnectionError(err_msg)
                                       
    def get_ip(self):
        return self.svr_ip

    def get_start_time(self):
        return self.svr_start_time
                                       
    def _generate_job_def(self, settings_path):
        from dexen.system.client.houdini_generator import HoudiniGenerator
        h_gen = HoudiniGenerator(settings_path)
        h_gen.generate()        

    def _is_valid_job_name(self, job_name):
        if not job_name or len(job_name) == 0: 
            return False
        for i in xrange(len(job_name)):
            if not job_name[i].isalnum() and not job_name[i] == "_":
                return False
        return True

    def start_job(self, job_name, job_def_dir):
        settings_path = os.path.join(job_def_dir, "settings.xml")
        master_task_path = os.path.join(job_def_dir, "master_task.py")
        
        # If the job definition directory is not a valid job def
        if not os.path.isfile(settings_path) and \
           not os.path.isfile(master_task_path):
            err_msg = "The directory for the job definition is not valid.\n"
            err_msg += "It must contain either master task (master_task.py)"
            err_msg += " or settings file (settings.xml)."
            raise DexenInvalidJobDefError(err_msg)
        
        if os.path.isfile(settings_path):
            self._generate_job_def(settings_path)

        if not self._is_valid_job_name(job_name):
            err_msg = "The specified job name is not valid.\n"
            err_msg += "The job name can only contain alphanumeric characters"
            err_msg += " or '_' character."
            raise DexenInvalidJobNameError(err_msg)

        job_def_zip = utils.get_schema_zip(job_def_dir) 
        self.srv_conn.root.start_job(job_def_zip, job_name)
        
    def get_job_states(self):
        job_states = self.srv_conn.root.get_job_states()
        job_states = rpyc.classic.obtain(job_states)
        return job_states
    
    def get_job_proxy(self, job_id):
        if not self.job_proxies.has_key(job_id):
            self.job_proxies[job_id] = JobProxy(job_id, self.srv_conn)
        return self.job_proxies[job_id]
    
    def get_job_old(self, job_id):
        jobs = self.get_job_states()
        job_ids = [job.job_id for job in jobs]
        if not job_id in job_ids:
            if self.job_proxies.has_key(job_id):
                del self.job_proxies[job_id]
            raise JobNotExistsException("Job with id:%s doesn't exist."%(job_id))
        if self.job_proxies.has_key(job_id):
            return self.job_proxies[job_id]
        self.job_proxies[job_id] = JobProxy(job_id, self.srv_conn)
        return self.job_proxies[job_id]
        

#===============================================================================
# Job Proxy
#===============================================================================
class JobProxy(object):
    def __init__(self, job_id, srv_conn):
        self.job_id = job_id
        self.actions = []
        self.inds = {} # id:Individual
        self.last_act_id = 0
        self.recent_acts = []
        self.p_ranking = ParetoRanking()
        self.scores = [] # List of ScoreMeta
        self.lock = threading.RLock()
        self.srv_conn = srv_conn
        if self.srv_conn:
            self._process_config_file()
    
    def reconnect(self, srv_conn):
        self.srv_conn = srv_conn
    
    def _process_config_file(self):
        config_file = self.srv_conn.root.get_client_config_file(self.job_id)
        if config_file != None:
            exec config_file
        variables = locals()
        if variables.has_key("scores"):
            scores = variables["scores"]
            if isinstance(scores, dict):
                for name, ranking_type in scores.items():
                    self._add_score(name, ranking_type)
    
    def _add_score(self, name, ranking_type):
        ranking_type = ranking_type.upper()
        print "Adding score", name, ranking_type
        if ranking_type == "MIN":
            self.p_ranking.register_score(name, ranking.MIN)
        elif ranking_type == "MAX": 
            self.p_ranking.register_score(name, ranking.MAX)
    
    def stop(self):
        self.srv_conn.root.stop_job(self.job_id)
    
    def delete(self):
        self.srv_conn.root.del_job(self.job_id)
       
    def _update(self):
        new_acts = self.srv_conn.root.update(self.job_id, self.last_act_id)
        new_acts = rpyc.classic.obtain(new_acts)
        if len(new_acts) > 0:
            self.process_actions(new_acts)
            self.recent_acts.extend(new_acts)
        return len(new_acts)
    
    def process_actions(self, new_actions):
        for action in new_actions:
            assert isinstance(action, Action)
            self.last_act_id = max(self.last_act_id, action.get_id()) 
            inds = action.get_inds()
            for ind in inds:
                assert isinstance(ind, Individual)
                self.inds[ind._get_id()] = ind
        self.actions.extend(new_actions)

    def refresh(self):
        self._update()
    
    def get_recent_actions(self):
        acts = self.recent_acts
        self.recent_acts = []
        return acts
        
    def get_all_actions(self):
        return self.actions
    
    def get_actions_deprecated(self, from_index):
        self._update()
        return self.actions[from_index:]
    
    def get_actions(self, last_act_id):
        new_acts = self.srv_conn.root.update(self.job_id, self.last_act_id)
        new_acts = rpyc.classic.obtain(new_acts)
        if len(new_acts) > 0:
            action = new_acts[-1]
            assert isinstance(action, Action)
            self.last_act_id = action.get_id()
            assert action.get_id() - last_act_id == len(new_acts)
        self.actions.extend(new_acts)
        self.recent_acts.extend(new_acts)
        self.process_actions(new_acts)
        return new_acts
    
    def get_inds(self):
        return self.inds.values()
    
    def get_action(self, _id):
        assert _id <= len(self.actions)
        return self.actions[_id - 1]
    
    def get_ind_archive_names(self, ind_id):
        archiveNames = self.srv_conn.root.get_ind_archive_names(self.job_id, ind_id)
        archiveNames = rpyc.classic.obtain(archiveNames)
        return archiveNames
    
    def get_ind_archive(self, ind_id, name):
        content = self.srv_conn.root.get_ind_archive(self.job_id, ind_id, name)
        content = rpyc.classic.obtain(content)
        return content
    
    def get_ind_archive_all(self, ind_id):
        archives = self.srv_conn.root.get_ind_archive_all(self.job_id, ind_id)
        archives = rpyc.classic.obtain(archives)
        return archives # {name: content}
    
    def get_job_info(self):
        job_state = self.srv_conn.root.get_job_info(self.job_id)
        job_state = rpyc.classic.obtain(job_state)
        return job_state
    
    def get_zip(self):
        zip_data = self.srv_conn.root.get_job_def(self.job_id)
        zip_data = rpyc.classic.obtain(zip_data)
        return zip_data
     
    def get_pareto_log(self, max_level=None):
        self._update()
        log = ""
        inds = self.inds.values()      
        self.p_ranking.set_inds(inds)
        scores = self.p_ranking.get_scores()
        inds = self.p_ranking.get_fully_evaluated_inds()
        
        header = Individual._get_header()
        for score_meta in scores:
            assert isinstance(score_meta, ScoreMeta)
            header += ",%s"%(score_meta.name)
        log += header + ",Rank\n"
                        
        for ind in inds:
            assert isinstance(ind, Individual)
            log += ind._get_pareto_log()
            for score_meta in scores:
                log += ",%f"%(getattr(ind, score_meta.name))
            log += ",None\n"
        return log   
    
    def get_scores(self):
        # Return list of ScoreMeta 
        return self.p_ranking.get_scores()
        
    def get_progress_log(self, step_size):
        self._update()
        scores = self.get_scores()
        log = "%s\n"%(ScoreProgress.get_header(scores),)
        score_progress = ScoreProgress(self.inds.values(), scores, step_size)
        for time, score_info in score_progress.get_progress():
            log += "%d"%(time,)
            for best, avg, worst in score_info.values():
                log += ",%f,%f,%f"%(best, avg, worst) 
            log += "\n"
        return log

    def get_advancement_log(self, step):
        self._update()
        scores = self.get_scores()
        p_adv = ParetoAdvancement(self.inds.values(), scores, step)
        adv_list = p_adv.get_advancement()
        log = "%s\n"%(ParetoAdvancement.get_header(scores, adv_list))      
        row = 0
        while True:
            processed = False
            for time, inds in adv_list:
                if row >= len(inds):
                    log += " ,"*(len(scores)+1)
                    continue                 
                processed = True
                ind = inds[row]
                assert isinstance(ind, Individual)
                log += "%d"%(ind.__id__)
                for score in scores:
                    assert isinstance(score, ScoreMeta)
                    log += ",%f"%(getattr(ind, score.name))
                log += ","
            log += "\n"
            if not processed: 
                break
            row += 1
        return log

    # Deprecated: only used by CLI
    # TODO: To be removed in the future
    def get_action_log(self):
        self._update()
        log = ""
        if len(self.actions) > 0:
            action = self.actions[-1]
            assert isinstance(action, Action)
            log += action.get_header() + "\n"
        
        for action in self.actions:
            assert isinstance(action, Action)
            log += action.get_log() + "\n"
        return log

    # Deprecated: only used by CLI
    # TODO: To be removed in the future
    def get_output_log(self):
        self._update()
        log = ""
        for action in self.actions:
            assert isinstance(action, Action)
            if action.is_receive():
                #transact = action.get_transact()
                #assert isinstance(transact, RecvTransact)
                #log += transact.__repr__() + "\n"
                log += action.__repr__() + "\n"     
        return log
    
    
